#!/bin/bash

# Copyright (c) 2020 Teradici Corporation
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

exec > >(tee /var/log/user-data.log|logger -t user-data -s 2>/dev/console) 2>&1

LOG_FILE="/var/log/teradici/provisioning.log"
METADATA_IP="http://169.254.169.254"
TERADICI_REPO_SETUP_SCRIPT_URL="https://dl.teradici.com/${teradici_download_token}/pcoip-agent/cfg/setup/bash.rpm.sh"

log() {
    local message="$1"
    echo "[$(date)] $message"
}

# Try command until zero exit status or exit(1) when non-zero status after max tries
retry() {
    local counter="$1"
    local interval="$2"
    local command="$3"
    local log_message="$4"
    local err_message="$5"
    local count=0

    while [ true ]
    do
        ((count=count+1))
        eval $command && break
        if [ $count -gt $counter ]
        then
            log "$err_message"
            return 1
        else
            log "$log_message Retrying in $interval seconds"
            sleep $interval
        fi
    done
}

get_credentials() {
    # Disable logging of secrets by wrapping the region with set +x and set -x
    set +x
    if [[ -z "${customer_master_key_id}" ]]; then
        log "--> Script is not using encryption for secrets."

        PCOIP_REGISTRATION_CODE="${pcoip_registration_code}"
        AD_SERVICE_ACCOUNT_PASSWORD="${ad_service_account_password}"

    else
        log "--> Script is using encryption key: '${customer_master_key_id}'."

        if [[ "${pcoip_registration_code}" ]]; then
            log "--> Decrypting pcoip_registration_code..."
            PCOIP_REGISTRATION_CODE=$(aws kms decrypt --region ${aws_region} --ciphertext-blob fileb://<(echo "${pcoip_registration_code}" | base64 -d) --output text --query Plaintext | base64 -d)
        fi

        log "--> Decrypting ad_service_account_password..."
        AD_SERVICE_ACCOUNT_PASSWORD=$(aws kms decrypt --region ${aws_region} --ciphertext-blob fileb://<(echo "${ad_service_account_password}" | base64 -d) --output text --query Plaintext | base64 -d)
    fi
    set -x
}

check_required_vars() {
    set +x
    if [[ -z "$AD_SERVICE_ACCOUNT_PASSWORD" ]]; then
        log "--> ERROR: Missing Active Directory Service Account Password."
        missing_vars="true"
    fi
    set -x

    if [[ "$missing_vars" = "true" ]]; then
        log "--> Exiting..."
        exit 1
    fi
}

# Update the hostname to match this instance's "Name" Tag
update_hostname() {
    TOKEN=`curl -X PUT "$METADATA_IP/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 60"`
    ID=`curl $METADATA_IP/latest/meta-data/instance-id -H "X-aws-ec2-metadata-token: $TOKEN"`
    REGION=`curl $METADATA_IP/latest/dynamic/instance-identity/document/ -H "X-aws-ec2-metadata-token: $TOKEN" | jq -r .region`
    NEW_HOSTNAME=`aws ec2 describe-tags --region $REGION --filters "Name=resource-id,Values=$ID" "Name=key,Values=Name" --output json | jq -r .Tags[0].Value`

    sudo hostnamectl set-hostname $NEW_HOSTNAME.${domain_name}
}

exit_and_restart() {
    log "--> Rebooting..."
    (sleep 1; reboot -p) &
    exit
}

install_kernel_header() {
    log "--> Installing kernel headers and development packages..."
    yum install kernel-devel kernel-headers -y
    exitCode=$?
    if [[ $exitCode -ne 0 ]]; then
        log "--> ERROR: Failed to install kernel header."
        exit 1
    fi

    yum install gcc -y
    exitCode=$?
    if [[ $exitCode -ne 0 ]]; then
        log "--> ERROR: Failed to install gcc."
        exit 1
    fi
}

remove_nouveau() {
        log "--> Disabling the Nouveau kernel driver..."
        for driver in vga16fb nouveau nvidiafb rivafb rivatv; do
            echo "blacklist $driver" >> /etc/modprobe.d/blacklist.conf
        done

        sed -i 's/\(^GRUB_CMDLINE_LINUX=".*\)"/\1 rdblacklist=nouveau"/' /etc/default/grub
        grub2-mkconfig -o /boot/grub2/grub.cfg
}

# Download installation script and run to install NVIDIA driver
install_gpu_driver() {
    # the first part to check if GPU is attached
    # NVIDIA VID = 10DE
    # Display class code = 0300
    # the second part to check if the NVIDIA driver is installed
    if [[ $(lspci -d '10de:*:0300' -s '.0' | wc -l) -gt 0 ]] && ! (modprobe --resolve-alias nvidia > /dev/null 2>&1)
    then
        log "--> Installing GPU driver..."

        local nvidia_driver_filename=$(basename ${nvidia_driver_url})
        local gpu_installer="/root/$nvidia_driver_filename"

        log "--> Killing X server before installing driver..."
        systemctl stop gdm

        log "--> Downloading GPU driver $nvidia_driver_filename to $gpu_installer..."
        retry   3 \
                5 \
                "curl -f -o $gpu_installer ${nvidia_driver_url}" \
                "--> Non-zero exit status." \
                "--> ERROR: Failed to download GPU driver installer."
        if [ $? -eq 1 ]; then
            exit 1
        fi
        log "--> GPU driver installer successfully downloaded"

        chmod u+x "$gpu_installer"
        log "--> Running GPU driver installer..."
        # -s, --silent Run silently; no questions are asked and no output is printed,
        # This option implies '--ui=none --no-questions'.
        # -X, --run-nvidia-xconfig
        # -Z, --disable-nouveau
        # --sanity Perform basic sanity tests on an existing NVIDIA driver installation.
        # --uninstall Uninstall the currently installed NVIDIA driver.
        # using dkms cause kernel rebuild and installation failure
        if "$gpu_installer" -s -Z -X
        then
            log "--> GPU driver installed successfully."
        else
            log "--> ERROR: Failed to install GPU driver."
            exit 1
        fi
    fi
}

# Enable persistence mode
enable_persistence_mode() {
    # the first part to check if the NVIDIA driver is installed
    # the second part to check if persistence mode is enabled
    if (modprobe --resolve-alias nvidia > /dev/null 2>&1) && [[ $(nvidia-smi -q | awk '/Persistence Mode/{print $NF}') != "Enabled" ]]
    then
        log "--> Enabling persistence mode..."

        # tar -xvjf /usr/share/doc/NVIDIA_GLX-1.0/sample/nvidia-persistenced-init.tar.bz2 -C /tmp
        # chmod +x /tmp/nvidia-persistenced-init/install.sh
        # /tmp/nvidia-persistenced-init/install.sh
        # local exitCode=$?
        # rm -rf /tmp/nvidia-persistenced-init
        # if [[ $exitCode -ne 0 ]]

        # Enable persistence mode
        # based on document https://docs.nvidia.com/deploy/driver-persistence/index.html,
        # Persistence Daemon shall be used in future
        if (nvidia-smi -pm ENABLED)
        then
            log "--> Persistence mode is enabled successfully"
        else
            log "--> ERROR: Failed to enable persistence mode."
            exit 1
        fi
    fi
}

install_pcoip_agent() {
    if ! (rpm -q pcoip-agent-graphics)
    then

        log "--> Getting Teradici PCoIP agent repo..."
        curl --retry 3 --retry-delay 5 -u "token:${teradici_download_token}" -1sLf $TERADICI_REPO_SETUP_SCRIPT_URL | bash
        if [ $? -ne 0 ]; then
            log "--> ERROR: Failed to install PCoIP agent repo."
            exit 1
        fi
        log "--> PCoIP agent repo installed successfully."

        log "--> Installing USB dependencies..."
        retry   3 \
                5 \
                "yum install -y usb-vhci" \
                "--> Non-zero exit status." \
                "--> Warning: Failed to install usb-vhci."

        log "--> Installing PCoIP graphics agent..."
        retry   3 \
                5 \
                "yum -y install pcoip-agent-graphics" \
                "--> Non-zero exit status." \
                "--> ERROR: Failed to download PCoIP agent."
        if [ $? -ne 0 ]; then
            exit 1
        fi
        log "--> PCoIP agent installed successfully."

        set +x
        if [[ "${pcoip_registration_code}" ]]; then
            log "--> Registering PCoIP agent license..."
            n=0
            while true; do
                /usr/sbin/pcoip-register-host --registration-code="$PCOIP_REGISTRATION_CODE" && break
                n=$[$n+1]

                if [ $n -ge 10 ]; then
                    log "--> ERROR: Failed to register PCoIP agent after $n tries."
                    exit 1
                fi

                log "--> ERROR: Failed to register PCoIP agent. Retrying in 10s..."
                sleep 10
            done
            log "--> PCoIP agent registered successfully."

        else
            log "--> No PCoIP Registration Code provided. Skipping PCoIP agent registration..."
        fi
        set -x
    fi
}

join_domain() {
    local dns_record_file="dns_record"
    if [[ ! -f "$dns_record_file" ]]
    then
        log "--> DOMAIN NAME: ${domain_name}"
        log "--> USERNAME: ${ad_service_account_username}"
        log "--> DOMAIN CONTROLLER: ${domain_controller_ip}"

        # default hostname has the form ip-10-0-0-1.us-west-1.compute.internal,
        # get the first part of it
        VM_NAME=$(echo $(hostname) | sed -n 's/\(^[^.]*\).*/\1/p')

        # Wait for AD service account to be set up
        yum -y install openldap-clients
        log "--> Waiting for AD account ${ad_service_account_username}@${domain_name} to be available..."
        set +x
        until ldapwhoami -H ldap://${domain_controller_ip} -D ${ad_service_account_username}@${domain_name} -w "$AD_SERVICE_ACCOUNT_PASSWORD" -o nettimeout=1 > /dev/null 2>&1
        do
            log "--> ${ad_service_account_username}@${domain_name} not available yet, retrying in 10 seconds..."
            sleep 10
        done
        set -x

        # Join domain
        log "--> Installing required packages to join domain..."
        yum -y install sssd realmd oddjob oddjob-mkhomedir adcli samba-common samba-common-tools krb5-workstation openldap-clients policycoreutils-python

        log "--> Joining the domain '${domain_name}'..."
        local retries=10

        set +x
        while true
        do
            echo "$AD_SERVICE_ACCOUNT_PASSWORD" | realm join --user="${ad_service_account_username}@${domain_name}" "${domain_name}" --verbose >&2

            local rc=$?
            if [[ $rc -eq 0 ]]
            then
                log "--> Successfully joined domain '${domain_name}'."
                break
            fi

            if [ $retries -eq 0 ]
            then
                log "--> ERROR: Failed to join domain '${domain_name}'."
                return 106
            fi

            log "--> ERROR: Failed to join domain '${domain_name}'. $retries retries remaining..."
            retries=$((retries-1))
            sleep 60
        done
        set -x

        domainname "$VM_NAME.${domain_name}"
        echo "%${domain_name}\\\\Domain\\ Admins ALL=(ALL) ALL" > /etc/sudoers.d/sudoers

        log "--> Registering with DNS..."
        DOMAIN_UPPER=$(echo "${domain_name}" | tr '[:lower:]' '[:upper:]')
        IP_ADDRESS=$(hostname -I | grep -Eo '10.([0-9]*\.){2}[0-9]*')
        set +x
        echo "$AD_SERVICE_ACCOUNT_PASSWORD" | kinit "${ad_service_account_username}"@"$DOMAIN_UPPER"
        set -x
        touch "$dns_record_file"
        echo "server ${domain_controller_ip}" > "$dns_record_file"
        echo "update add $VM_NAME.${domain_name} 600 a $IP_ADDRESS" >> "$dns_record_file"
        echo "send" >> "$dns_record_file"
        nsupdate -g "$dns_record_file" > /var/sky

        log "--> Configuring settings..."
        sed -i '$ a\dyndns_update = True\ndyndns_ttl = 3600\ndyndns_refresh_interval = 43200\ndyndns_update_ptr = True\nldap_user_principal = nosuchattribute' /etc/sssd/sssd.conf
        sed -c -i "s/\\(use_fully_qualified_names *= *\\).*/\\1False/" /etc/sssd/sssd.conf
        sed -c -i "s/\\(fallback_homedir *= *\\).*/\\1\\/home\\/%u/" /etc/sssd/sssd.conf

        # sssd.conf configuration is required first before enabling sssd
        log "--> Restarting messagebus service..."
        if ! (systemctl restart messagebus)
        then
            log "--> ERROR: Failed to restart messagebus service."
            return 106
        fi

        log "--> Enabling and starting sssd service..."
        if ! (systemctl enable sssd --now)
        then
            log "--> ERROR: Failed to start sssd service."
            return 106
        fi
    fi
}

# Open up firewall for PCoIP Agent. By default eth0 is in firewall zone "public"
update_firewall() {
    log "--> Adding 'pcoip-agent' service to public firewall zone..."
    firewall-offline-cmd --zone=public --add-service=pcoip-agent
    systemctl enable firewalld
    systemctl start firewalld
}

# A flag to indicate if this is run from reboot
RE_ENTER=0

if (rpm -q pcoip-agent-graphics); then
    exit
fi

if [[ ! -f "$LOG_FILE" ]]
then
    mkdir -p "$(dirname $LOG_FILE)"
    touch "$LOG_FILE"
    chmod +644 "$LOG_FILE"
else
    RE_ENTER=1
fi

log "$(date)"

# Print all executed commands to the terminal
set -x

# Redirect stdout and stderr to the log file
exec &>>$LOG_FILE

get_credentials

check_required_vars

if [[ $RE_ENTER -eq 0 ]]
then
    # EPEL needed for GraphicsMagick-c++, required by PCoIP Agent
    yum -y install epel-release
    yum -y update
    yum install -y wget awscli jq

    update_hostname

    # Install GNOME and set it as the desktop
    log "--> Installing Linux GUI..."
    yum -y groupinstall "GNOME Desktop" "Graphical Administration Tools"

    log "--> Setting default to graphical target..."
    systemctl set-default graphical.target

    join_domain

    remove_nouveau

    exit_and_restart
else
    install_kernel_header

    install_gpu_driver

    enable_persistence_mode

    install_pcoip_agent

    update_firewall

    log "--> Installation is complete!"

    exit_and_restart
fi
